#!/usr/bin/env python
# encoding: utf-8

import os
import argparse
import numpy as np
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA
from sklearn.decomposition import PCA
from sklearn import preprocessing

def compute_eer(target_scores, nontarget_scores):
    if isinstance(target_scores , list) is False:
        target_scores = list(target_scores)
    if isinstance(nontarget_scores , list) is False:
        nontarget_scores = list(nontarget_scores)

    target_scores = sorted(target_scores)
    nontarget_scores = sorted(nontarget_scores)
    target_size = len(target_scores);
    nontarget_size = len(nontarget_scores)

    target_position = 0
    for i in range(target_size-1):
        target_position = i
        nontarget_n = nontarget_size * float(target_position) / target_size
        nontarget_position = int(nontarget_size - 1 - nontarget_n)
        if nontarget_position < 0:
            nontarget_position = 0
        if nontarget_scores[nontarget_position] < target_scores[target_position]:
            break
    th = target_scores[target_position]
    eer = target_position * 1.0 / target_size
    return eer, th


def load_npy(npy_dir):
    print("loading npy data from {}".format(npy_dir))
    spks = os.listdir(npy_dir)
    labels = []
    vectors = []
    for spk in spks:
        dirname = os.path.join(npy_dir, spk)
        for npy in os.listdir(dirname):
            path = os.path.join(dirname, npy)
            vectors.append(np.load(path))
            labels.append(spk)
    vectors = np.array(vectors)
    return vectors, labels


def cosine_score(enroll_vectors, enroll_labels, test_vectors, test_labels):
    cosine_matrix = enroll_vectors.dot(test_vectors.T)
    enroll_norm = np.linalg.norm(enroll_vectors, axis=1, keepdims=True)
    test_norm = np.linalg.norm(test_vectors, axis=1, keepdims=True)
    denom = enroll_vectors.dot(test_vectors.T)
    cosine_matrix = cosine_matrix/enroll_norm.dot(test_norm.T)

    target_scores = []
    nontarget_scores = []
    for x in range(len(enroll_labels)):
        for y in range(len(test_labels)):
            if enroll_labels[x] == test_labels[y]:
                target_scores.append(cosine_matrix[x][y])
            else:
                nontarget_scores.append(cosine_matrix[x][y])
    eer, th = compute_eer(target_scores, nontarget_scores)
    return eer, th


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--dev_dir', help='data list', type=str, default="data/tmp_list")
    parser.add_argument('--enroll_dir', help='data list', type=str, default="data/tmp_list")
    parser.add_argument('--test_dir', help='embedding dim', type=str, default="data/test_list")
    args = parser.parse_args()

    dev_vectors, dev_labels = load_npy(args.dev_dir)
    enroll_vectors, enroll_labels = load_npy(args.enroll_dir)
    test_vectors, test_labels = load_npy(args.test_dir)

    embedding_dim = dev_vectors.shape[-1]
    print("embedding dim is {}".format(embedding_dim))
    eer, _ = cosine_score(enroll_vectors, enroll_labels, test_vectors, test_labels)
    print("Cosine EER: {:.3f}%".format(eer*100))

    for dim in [16, 32, 64, 128, 150, 256]:
        pca = PCA(n_components=dim, whiten=False)
        dev_vectors_trans = pca.fit_transform(dev_vectors)
        enroll_vectors_trans = pca.transform(enroll_vectors)
        test_vectors_trans = pca.transform(test_vectors)
        eer, _ = cosine_score(enroll_vectors_trans, enroll_labels, test_vectors_trans, test_labels)
        print("PCA {} Cosine EER: {:.3f}%".format(dim, eer*100))

    for dim in [8, 16, 24, 32, 40, 48, 49]:
        lda = LDA(n_components=dim)
        dev_vectors_trans = lda.fit_transform(dev_vectors, dev_labels)
        enroll_vectors_trans = lda.transform(enroll_vectors)
        test_vectors_trans = lda.transform(test_vectors)
        eer, _ = cosine_score(enroll_vectors_trans, enroll_labels, test_vectors_trans, test_labels)
        print("LDA {} Cosine EER: {:.3f}%".format(dim, eer*100))

